#Data : 2025-3-20
#Author : Fengyuan Zhang (Franklin)
#Email : franklinzhang@foxmail.com
#Description : Demo for model encapsulation (YOLOv5)

import torch
import cv2
import numpy as np
import sys, os
from modelservicecontext import EModelContextStatus
from modelservicecontext import ERequestResponseDataFlag
from modelservicecontext import ERequestResponseDataMIME
from modelservicecontextlite import ModelServiceContextLite
from modeldatahandler import ModelDataHandler

ms = ModelServiceContextLite.createModelServiceContext(sys.argv)
if (ms == None):
    exit()

ms.onInitialize()

db = ms.getRequestData("Run", "LoadImage")
if db is not None:
    df = db.getByFile()
else:
    ms.onFinalize()

# model = torch.hub.load(os.path.dirname(__file__) + '/yolov5', 'custom', os.path.dirname(__file__) + "/yolov5s.pt", source='local', weights_only=False)
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

img = cv2.imread(df)
if img is None:
    ms.onPostErrorInfo(f"Can not read the image : {df}")
    raise FileNotFoundError(f"Can not read the image : {df}")

# YOLOv5 模型要求输入为 RGB 图像，因此需要转换 BGR -> RGB
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

# 3. 使用模型进行目标检测
results = model(img_rgb)

# 4. 打印检测结果（包含检测框、置信度、类别等信息）

result1 = results.pandas().xyxy[0]
# print(results.pandas().xyxy[0])
resultfile1 = ms.getCurrentDataDirectory() + "result.dat"
f = open(resultfile1, "w+")
f.write(str(result1))
f.close()
ms.setResponseDataByStream("Run", "ResultData", resultfile1)

# 5. 显示检测结果（会自动绘制检测框并展示图像）
resultfile2 = ms.getCurrentDataDirectory(False)[:-1]
# results.show()
# cv2.imwrite(result2, results)
results.save(save_dir = resultfile2)
resultfile2 = resultfile2 + ms.getSlash() + "image0.jpg"
# results.save(result2)
ms.setResponseDataByFile("Run", "ResultImage", resultfile2)

ms.onFinalize()

# 6. 若需要以 numpy 数组形式获取带检测框的图像，可以使用：
# annotated_img = np.squeeze(results.render())[:, :, ::-1]  # 转换回 BGR 格式
# cv2.imshow("Detection", annotated_img)
# cv2.waitKey(0)+
# cv2.destroyAllWindows()